import numpy as np
from matplotlib import pyplot as plt
from scipy import optimize 
from scipy.optimize import curve_fit
from matplotlib.ticker import (MultipleLocator,FormatStrFormatter,AutoMinorLocator)

exp = np.loadtxt('600K_exp_param.txt',skiprows=1)
exp_ce = np.loadtxt('600K_cefit.txt')

exp_HWHM=exp[:,13]/2.0
exp_ERROR=np.sqrt(exp[:,14])/2

##########################
Title=r'$MLMD$'
#Title=r'$7Li_6PS_5Cl$'
Tit='7Li6PS5Cl'
fs=20
ymin=0
#ymax=input("Ymax\n")
#ymax=int(ymax)
ymax=150
xmin=0.0
xmax=2.0
xmajor=0.5
xminor=0.1
ymajor=50
yminor=10

p0=[0.1,2.0]
###########FOR Result Printing
ypos=0.92*ymax
dely=0.1*ymax
xpos=0.08

###########################
#TEMP=input("Number of Lorentzian\n")
TEMP = 2
#filename="%s_CE_MODEL600K%sLor_slow_narrow.pdf"%(Tit,TEMP)
filename="%s_%s_Exp_MLMD.pdf"%(Tit,TEMP)
exp_filename="600K_param_%sL.txt"%TEMP
microtops=1517
a=np.loadtxt(exp_filename,skiprows=1)
ENE=a[:,0]
ENE2=(np.arange(xmin,xmax,0.1))
#######################FWHM TO HWHM 
HWHM=a[:,1]/2.0
ERROR=np.sqrt(a[:,1])/2
##################################


def CE(qval,tjump,djump):
    arg=qval*djump
    y= ( 1 - (np.sin(arg)/arg) )/tjump
    return y

best_vals, covar = curve_fit(CE, ENE, HWHM, p0=p0, bounds=([0,0],[1000,4.0]))
perr = np.sqrt(np.diag(covar))
tjump,djump=best_vals
terr,derr=perr #best_vals
print(perr)
#print(tjump,djump)
yfit=CE(ENE2,best_vals[0],best_vals[1])
fig=plt.figure(figsize=(8,6))
ax=fig.add_subplot()

ax.plot(exp_ce[:,0],exp_ce[:,1],color='C0')
ax.errorbar(exp[:,0],exp_HWHM,exp_ERROR,fmt='o',color='C0')
exp = np.column_stack((exp[:,0],exp_HWHM,exp_ERROR))
np.savetxt('exp_CEfit',exp)

ax.plot(ENE2,yfit,color='C1')
ax.errorbar(ENE,HWHM,ERROR,fmt='o',color='C1')
ax.set_xlabel("Q($\AA^{-1}$)",fontsize=fs)

MLMD = np.column_stack((ENE,HWHM,ERROR))
np.savetxt('MLMD_CEfit',MLMD)
#ax.set_ylabel(r"$\Gamma$(\u03bceV)",fontsize=fs)
#ax.set_ylabel(r"$\Gamma$/2($\mu$eV)",fontsize=fs)
ax.set_ylabel(r"$\Gamma$($\mu$eV)",fontsize=fs)
tpicosec=tjump*microtops
terr=terr*microtops
Diffusion=djump*djump/(6.0*tpicosec)
Differr=Diffusion*(2*derr/djump + terr/tpicosec)
Differr=Differr*100
print(derr,terr,Differr)
Diffusion=Diffusion*100  #to make it 10^-10 m2/sec (10^-6 cm2/sec)


#ymax=np.max(HWHM)
#print(xpos,ypos)

ax.set_xticks(ENE2[0::])
#ax.set_xlim(xmin,xmax)
ax.set_xlim(xmin,2)
ax.set_ylim(ymin,ymax)
#ax.text(1.0,0.9*ymax,"%s"%Title,fontsize=fs)
ax.text(1.5,ypos,"T=600K",fontsize=fs)

'''
text_d = [r'$d_{jump}$=',r'2.8$\pm$0.1',';',r'2.4$\pm$0.4',r' $\AA$']
Color=['black','C1','black','C0','black']
r = fig.canvas.get_renderer()

space = 0
for i in range(5):
    t = ax.text(xpos,ypos,text_d[i],color=Color[i],fontsize=fs,ha='left')
    transf = ax.transData.inverted()
    bb = t.get_window_extent(renderer=fig.canvas.renderer)
    bb = bb.transformed(transf)
    xpos = xpos + bb.xmax-bb.xmin + space
xpos=0.1

text_d = [r'$\tau$=',r'18.7$\pm$0.4',';',r'19.2$\pm$2.4',' ps']

space = 0
for i in range(5):
    t = ax.text(xpos,ypos-dely,text_d[i],color=Color[i],fontsize=fs,ha='left')
    transf = ax.transData.inverted()
    bb = t.get_window_extent(renderer=fig.canvas.renderer)
    bb = bb.transformed(transf)
    xpos = xpos + bb.xmax-bb.xmin + space
xpos=0.1

text_d = [r'D=',r'(6.9$\pm$0.5)',';',r'(5.0$\pm$2.1)','x10$^{-6}$ cm$^2$/s']

space = 0
for i in range(5):
    t = ax.text(xpos,ypos-2*dely,text_d[i],color=Color[i],fontsize=fs,ha='left')
    transf = ax.transData.inverted()
    bb = t.get_window_extent(renderer=fig.canvas.renderer)
    bb = bb.transformed(transf)
    xpos = xpos + bb.xmax-bb.xmin + space
xpos=0.1
'''

ax.text(0.3,50,'$MLMD$',fontsize=fs,ha='left',color='C1')
ax.text(1.3,50,'$BASIS$',fontsize=fs,ha='left',color='C0')

#ax.text(xpos,ypos-dely,"$d_{jump}$=%3.1f $\pm$%3.1f   $\AA$"%(djump,derr),fontsize=fs)
#ax.text(xpos,ypos-2*dely,r"$\tau$=%3.1f $\pm$%3.1f ps"%(tpicosec,terr),fontsize=fs)
#ax.text(xpos,ypos-3*dely,"D=(%3.1f$\pm$%3.1f)x$10^{-6}  cm^2/sec$" %(Diffusion,Differr),fontsize=fs)

ax.xaxis.set_major_locator(MultipleLocator(xmajor))
ax.xaxis.set_major_formatter(FormatStrFormatter('%2.1f'))
ax.xaxis.set_minor_locator(MultipleLocator(xminor))
ax.yaxis.set_major_locator(MultipleLocator(ymajor))
ax.yaxis.set_minor_locator(MultipleLocator(yminor))
ax.tick_params(which='both',labelsize=20)
ax.tick_params(which='both', width=1)
ax.tick_params(which='major', length=8)
ax.tick_params(which='minor', length=4)

plt.savefig('QENS_exp_mlmd2.pdf',bbox_inches='tight',format='pdf')
plt.show()
